"""YOLOv8模型封装模块"""

import os
import cv2
import numpy as np
from typing import List, Dict, Tuple, Optional
from ultralytics import YOLO
from config import get_config

class YOLOv8Model:
    """YOLOv8模型封装类，提供统一的检测接口"""
    
    def __init__(self, model_path: str = None, device: str = 'auto'):
        """初始化YOLOv8模型
        
        Args:
            model_path: 模型文件路径，如果为None则使用配置中的路径
            device: 运行设备，'auto'、'cuda'、'cpu'
        """
        # 如果未指定模型路径，从配置中获取
        if model_path is None:
            model_path = get_config('detection.model_path', 'models/yolov8n.pt')
            
        # 确保模型路径存在
        if not os.path.isabs(model_path):
            base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
            model_path = os.path.join(base_dir, model_path)
        
        # 加载模型
        self.model = YOLO(model_path)
        
        # 设置设备
        self.device = device
        
        # 加载配置参数
        self.confidence_threshold = get_config('detection.confidence_threshold', 0.5)
        self.iou_threshold = get_config('detection.iou_threshold', 0.45)
        self.max_det = get_config('detection.max_det', 1000)
        
        # 获取类别信息
        self.classes = get_config('detection.classes', [])
        self.class_names = {cls['id']: cls['name'] for cls in self.classes}
        
        # 性能优化设置
        self.fp16 = get_config('performance.fp16', True)
        
        print(f"YOLOv8模型已加载: {model_path}")
    
    def detect(self, image: np.ndarray, classes: List[int] = None) -> List[Dict]:
        """执行目标检测
        
        Args:
            image: 输入图像 (BGR格式)
            classes: 需要检测的类别ID列表，None表示检测所有类别
            
        Returns:
            检测结果列表，每个元素包含: 
            {x1, y1, x2, y2, confidence, class_id, class_name}
        """
        try:
            # 执行检测
            results = self.model(
                image,
                conf=self.confidence_threshold,
                iou=self.iou_threshold,
                max_det=self.max_det,
                classes=classes,
                device=self.device,
                half=self.fp16,
                verbose=False
            )
            
            # 处理检测结果
            detections = []
            for result in results:
                boxes = result.boxes
                for box in boxes:
                    # 获取边界框坐标
                    x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
                    
                    # 获取置信度和类别ID
                    confidence = float(box.conf[0])
                    class_id = int(box.cls[0])
                    
                    # 获取类别名称
                    class_name = self.class_names.get(class_id, f'class_{class_id}')
                    
                    # 添加到检测结果列表
                    detections.append({
                        'x1': x1,
                        'y1': y1,
                        'x2': x2,
                        'y2': y2,
                        'confidence': confidence,
                        'class_id': class_id,
                        'class_name': class_name
                    })
            
            return detections
            
        except Exception as e:
            print(f"检测过程中发生错误: {e}")
            return []
    
    def detect_image_file(self, image_path: str, classes: List[int] = None) -> List[Dict]:
        """从图像文件执行检测
        
        Args:
            image_path: 图像文件路径
            classes: 需要检测的类别ID列表
            
        Returns:
            检测结果列表
        """
        try:
            # 读取图像
            image = cv2.imread(image_path)
            if image is None:
                raise ValueError(f"无法读取图像文件: {image_path}")
            
            # 执行检测
            return self.detect(image, classes)
            
        except Exception as e:
            print(f"图像检测失败: {e}")
            return []
    
    def draw_detections(self, image: np.ndarray, detections: List[Dict], 
                       show_confidence: bool = True, show_class: bool = True) -> np.ndarray:
        """在图像上绘制检测结果
        
        Args:
            image: 输入图像
            detections: 检测结果列表
            show_confidence: 是否显示置信度
            show_class: 是否显示类别名称
            
        Returns:
            绘制后的图像
        """
        # 创建图像副本以避免修改原图
        result_img = image.copy()
        
        # 为不同类别定义不同颜色
        colors = {
            'person': (0, 0, 255),    # 红色
            'car': (0, 255, 0),      # 绿色
            'bicycle': (255, 165, 0), # 橙色
            'motorcycle': (255, 0, 255), # 紫色
            'truck': (0, 255, 255)    # 青色
        }
        
        # 默认颜色
        default_color = (128, 128, 128)  # 灰色
        
        # 绘制每个检测结果
        for det in detections:
            # 获取坐标和类别信息
            x1, y1, x2, y2 = det['x1'], det['y1'], det['x2'], det['y2']
            class_name = det['class_name']
            confidence = det['confidence']
            
            # 获取颜色
            color = colors.get(class_name, default_color)
            
            # 绘制边界框
            cv2.rectangle(result_img, (x1, y1), (x2, y2), color, 2)
            
            # 准备标签文本
            label_parts = []
            if show_class:
                label_parts.append(class_name)
            if show_confidence:
                label_parts.append(f"{confidence:.2f}")
            
            if label_parts:
                label = ' '.join(label_parts)
                
                # 计算标签位置
                label_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)
                label_y = max(y1 - 10, label_size[1])
                
                # 绘制标签背景
                cv2.rectangle(
                    result_img,
                    (x1, label_y - label_size[1] - 5),
                    (x1 + label_size[0] + 10, label_y + 5),
                    color,
                    -1
                )
                
                # 绘制标签文本
                cv2.putText(
                    result_img,
                    label,
                    (x1 + 5, label_y),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.5,
                    (0, 0, 0),
                    2
                )
        
        return result_img
    
    def set_confidence_threshold(self, threshold: float) -> None:
        """设置置信度阈值
        
        Args:
            threshold: 置信度阈值 (0.0-1.0)
        """
        if 0.0 <= threshold <= 1.0:
            self.confidence_threshold = threshold
            print(f"置信度阈值已设置为: {threshold}")
        else:
            print("置信度阈值必须在0.0-1.0之间")
    
    def get_performance_metrics(self) -> Dict[str, any]:
        """获取模型性能指标
        
        Returns:
            性能指标字典
        """
        return {
            'model_name': self.model.model.__class__.__name__,
            'device': self.device,
            'confidence_threshold': self.confidence_threshold,
            'iou_threshold': self.iou_threshold,
            'max_detections': self.max_det,
            'fp16_enabled': self.fp16
        }
    
    def is_gpu_available(self) -> bool:
        """检查是否使用GPU
        
        Returns:
            是否使用GPU
        """
        try:
            import torch
            return torch.cuda.is_available() and 'cuda' in str(self.device).lower()
        except:
            return False